# -*- coding: utf-8 -*-
"""
基于akShare获取财务数据
"""
import akshare as ak
import pandas as pd
import pymysql
import requests
from bs4 import BeautifulSoup
from sqlalchemy import create_engine, text, Engine


def get_mysql_engine() -> Engine:
    """
    获取mysql数据库engine
    :return:
    """
    return create_engine('mysql+pymysql://root:123456@127.0.0.1:3308/stock?charset=utf8')


def add_income_statement_to_db():
    """
    ！！温馨提示：该方法及其耗时(我全量跑了大概5个小时)，近5000只股票循环调ak接口，然后ak接口里又是循环请求东财网页获取每一份财报，字段还多。

    把各个股票财务报表中的利润表主要数据存入数据库。

    因为获取财报数据接口较慢，而且每次只能查单只股票，而财报数据更新频率不高，相对较固定，存入数据库，需要的时候直接查库

    :return:
    """
    # 先获取所有股票，再逐个获取财报-利润表数据
    all_stock_list = list(ak.stock_zh_a_spot_em()["代码"])
    engine = get_mysql_engine()

    # 因为该方法及其耗时，所以加个已存在的判断，方便中途断了后不用重新来一遍
    db_exist_df = pd.read_sql(text("select distinct code from stock_income_statement"), engine.connect())
    db_exist_code = set() if db_exist_df.empty else set(db_exist_df["code"])

    # ak接口返回的column为大写，开始还费劲巴拉的挨个转为小写后再to_sql，跑完了才意识到直接大写to_sql应该也是可以的，一试果然可以。。。。
    db_columns = ["CODE", "NAME", "REPORT_DATE", "REPORT_NAME",  # 股票代码、报告日期
                  "OPERATE_INCOME", "OPERATE_INCOME_YOY",  # 营业收入
                  "TOTAL_PROFIT", "TOTAL_PROFIT_YOY",  # 利润总额
                  "NETPROFIT", "NETPROFIT_YOY",  # 净利润
                  "DEDUCT_PARENT_NETPROFIT", "DEDUCT_PARENT_NETPROFIT_YOY",  # 扣非净利润
                  "BASIC_EPS", "BASIC_EPS_YOY"]  # 每股收益
    i = 0
    count = len(all_stock_list)
    for stock in all_stock_list:
        i = i + 1
        print("%s  done......%d/%d" % (stock, i, count))
        if stock in db_exist_code:
            continue
        # 获取财报的接口，代码前面需要拼接"sh"/"sz"(沪市/深市)
        if stock.startswith("60") or stock.startswith("68"):
            stock = "sh" + stock
        elif stock.startswith("00") or stock.startswith("30"):
            stock = "sz" + stock
        else:
            continue
        # 获取财务报告-利润表数据（三季报数据为1-3季度汇总）
        # 本来想再获取单季度数据（stock_profit_sheet_by_quarterly_em），然后merge合成一条数据一起存库里，方便算TTM
        # 但这获取财务报表的接口本身就慢，再加单个季度的数据，执行更慢了，而且跑了一部分发现单季度数据很多缺的
        # 拉倒吧，单季度数据也能通过两个财报数据减出来
        try:
            # 如果没有财报数据，比如退市的，该接口会报错，try catch一下
            report_df = ak.stock_profit_sheet_by_report_em(stock)
        except KeyError:
            print(stock, "====调用ak接口报错了")
            continue

        # 股票代码
        report_df["CODE"] = report_df["SECURITY_CODE"]
        # 股票名称
        report_df["NAME"] = report_df["SECURITY_NAME_ABBR"]
        # 报告名称，2022三季报
        report_df["REPORT_NAME"] = report_df["REPORT_DATE_NAME"]

        # save to mysql
        db_df = report_df[db_columns]
        # 加上报表类型report_type: 1-一季报，2-中报，3-三季报，4-年报
        db_df.loc[db_df["REPORT_NAME"].str.contains('一季报'), "report_type"] = 1
        db_df.loc[db_df["REPORT_NAME"].str.contains('中报'), "report_type"] = 2
        db_df.loc[db_df["REPORT_NAME"].str.contains('三季报'), "report_type"] = 3
        db_df.loc[db_df["REPORT_NAME"].str.contains('年报'), "report_type"] = 4
        db_df.to_sql("stock_income_statement", engine, index=False, if_exists='append')
    # 关闭数据库连接
    engine.connect().close()


def add_new_income_statement_to_db():
    """
    增量获取新的财务报表数据（新发布尚未拉取存库的）

    :return:
    """
    # 获取已存在的最新
    query_sql = "select t.code, t.report_date from " \
                "(select *, rank() over(partition by code order by report_date desc) as time_rank " \
                "from `stock_income_statement`) t where t.time_rank = 1"
    engine = get_mysql_engine()
    exist_df = pd.read_sql(text(query_sql), engine.connect())
    exist_df["report_date"] = exist_df["report_date"].astype(str)
    engine.connect().close()
    if not exist_df.empty:
        # exist_date_dict["000001"] = "2022-12-31"
        exist_date_dict = exist_df.set_index("code").to_dict(orient='dict')["report_date"]

    # 先获取所有股票，再逐个获取财报-利润表数据
    all_stock_list = list(ak.stock_zh_a_spot_em()["代码"])
    engine = get_mysql_engine()

    # ak接口返回的column为大写，开始还费劲巴拉的挨个转为小写后再to_sql，跑完了才意识到直接大写to_sql应该也是可以的，一试果然可以。。。。
    db_columns = ["CODE", "NAME", "REPORT_DATE", "REPORT_NAME",  # 股票代码、报告日期
                  "OPERATE_INCOME", "OPERATE_INCOME_YOY",  # 营业收入
                  "TOTAL_PROFIT", "TOTAL_PROFIT_YOY",  # 利润总额
                  "NETPROFIT", "NETPROFIT_YOY",  # 净利润
                  "DEDUCT_PARENT_NETPROFIT", "DEDUCT_PARENT_NETPROFIT_YOY",  # 扣非净利润
                  "BASIC_EPS", "BASIC_EPS_YOY"]  # 每股收益
    i = 0
    count = len(all_stock_list)
    for stock in all_stock_list:
        i = i + 1
        print("%s  done......%d/%d" % (stock, i, count))
        # 获取已经存库的最新报告日期
        if stock in exist_date_dict:
            exist_last_date = exist_date_dict[stock]
        # 获取财报的接口，代码前面需要拼接"sh"/"sz"(沪市/深市)
        if stock.startswith("60") or stock.startswith("68"):
            stock = "sh" + stock
        elif stock.startswith("00") or stock.startswith("30"):
            stock = "sz" + stock
        else:
            continue
        # 获取财务报告-利润表数据（三季报数据为1-3季度汇总）
        # 本来想再获取单季度数据（stock_profit_sheet_by_quarterly_em），然后merge合成一条数据一起存库里，方便算TTM
        # 但这获取财务报表的接口本身就慢，再加单个季度的数据，执行更慢了，而且跑了一部分发现单季度数据很多缺的
        # 拉倒吧，单季度数据也能通过两个财报数据减出来
        try:
            # 如果没有财报数据，比如退市的，该接口会报错，try catch一下
            report_df = __new_profit_sheet_by_report(stock, exist_last_date)
            # 有可能没有新数据
            if report_df.empty:
                continue
        except KeyError:
            print(stock, "====调用ak接口报错了")
            continue

        # 股票代码
        report_df["CODE"] = report_df["SECURITY_CODE"]
        # 股票名称
        report_df["NAME"] = report_df["SECURITY_NAME_ABBR"]
        # 报告名称，2022三季报
        report_df["REPORT_NAME"] = report_df["REPORT_DATE_NAME"]

        # save to mysql
        db_df = report_df[db_columns].copy()
        # 加上报表类型report_type: 1-一季报，2-中报，3-三季报，4-年报
        db_df.loc[db_df["REPORT_NAME"].str.contains('一季报'), "report_type"] = 1
        db_df.loc[db_df["REPORT_NAME"].str.contains('中报'), "report_type"] = 2
        db_df.loc[db_df["REPORT_NAME"].str.contains('三季报'), "report_type"] = 3
        db_df.loc[db_df["REPORT_NAME"].str.contains('年报'), "report_type"] = 4
        db_df.to_sql("stock_income_statement", engine, index=False, if_exists='append')
        # print(db_df)
    # 关闭数据库连接
    engine.connect().close()


def __new_profit_sheet_by_report(symbol: str, last_report_date: str) -> pd.DataFrame:
    """
    获取指定日期之后新公布的财报-利润表

    参考ak.stock_profit_sheet_by_report_em，ak中只有获取一只票全部财报的接口，参考里面的改改获取指定报告日期之后的

    :param symbol: 股票代码; 带市场标识
    :param last_report_date 只获取报告日期在此之后的财报利润表，"2022-12-31"
    :return: 新公告的财报利润表
    """
    url = f"https://emweb.securities.eastmoney.com/PC_HSF10/NewFinanceAnalysis/Index"
    params = {"type": "web", "code": symbol.lower()}
    r = requests.get(url, params=params)
    soup = BeautifulSoup(r.text, "lxml")
    company_type = soup.find(attrs={"id": "hidctype"})["value"]
    url = "https://emweb.securities.eastmoney.com/PC_HSF10/NewFinanceAnalysis/lrbDateAjaxNew"
    params = {
        "companyType": company_type,
        "reportDateType": "0",
        "code": symbol,
    }
    r = requests.get(url, params=params)
    data_json = r.json()
    temp_df = pd.DataFrame(data_json["data"])
    temp_df["REPORT_DATE"] = pd.to_datetime(temp_df["REPORT_DATE"]).dt.date
    temp_df["REPORT_DATE"] = temp_df["REPORT_DATE"].astype(str)
    need_date = temp_df["REPORT_DATE"].tolist()
    date_list = [query_date for query_date in need_date if query_date > last_report_date]
    sep_list = [
        ",".join(date_list[i: i + 5]) for i in range(0, len(date_list), 5)
    ]
    big_df = pd.DataFrame()
    for item in sep_list:
        url = "https://emweb.securities.eastmoney.com/PC_HSF10/NewFinanceAnalysis/lrbAjaxNew"
        params = {
            "companyType": company_type,
            "reportDateType": "0",
            "reportType": "1",
            "code": symbol,
            "dates": item,
        }
        r = requests.get(url, params=params)
        data_json = r.json()
        temp_df = pd.DataFrame(data_json["data"])
        big_df = pd.concat([big_df, temp_df], ignore_index=True)
    return big_df


def get_income_statement_cagr(n_years=5, factors=[]) -> pd.DataFrame:
    """
    获取财务报表-利润表主要指标（营业收入、利润总额、净利润、扣非净利润、每股收益）在最近n年的【复合增长率（CAGR）】。

    基于已公布且存库的年报计算，可能存在滞后性（比如现在2023年3月，但一些公司尚未公布2022年报，所以只会计算截止到2021年的符合增长率）。

    ？？todo 如果期初、期末数据存在负数，返回结果为空，还不知道这个在财务上是怎么处理的。

    :param n_years:最近n年的复合增长率，默认5年，年报数据不足n年的股票会被忽略
    :param factors: 需要计算的指标，可选：[operate_income, total_profit, netprofit, deduct_parent_netprofit, basic_eps]
    :return:
    """
    query_sql = "select t.* from " \
                "(select *, rank() over(partition by code order by report_date desc) as time_rank " \
                "from `stock_income_statement` where report_type = 4) t " \
                "where t.time_rank in (1, %d) order by t.code, t.report_date desc" % n_years
    engine = get_mysql_engine()
    income_statement_df = pd.read_sql(text(query_sql), engine.connect())
    engine.connect().close()

    # 营业收入、利润总额、净利润、扣非净利润、每股收益
    if len(factors) == 0:
        columns = ["operate_income", "total_profit", "netprofit", "deduct_parent_netprofit", "basic_eps"]
    else:
        columns = factors
    # 把期初值及期末值放到一行，方便计算
    for column in columns:
        # 数据截止年份
        income_statement_df["end_year"] = income_statement_df["report_name"].shift(1).str.slice(0, 4)
        income_statement_df["latest_" + column] = income_statement_df[column].shift(1)
    # 只保留年报数量满足条件的公司
    income_statement_df = income_statement_df[income_statement_df["time_rank"] == n_years]
    return_columns = ["code", "name", "end_year"]
    # 期初值 * (1 + 复合增长率) ^ (n_years-1) = 期末值，即 复合增长率 = ((n_years-1)根号(期末值 / 期初值)) - 1
    for column in columns:
        return_columns.append(column + "_cagr")
        income_statement_df[column + "_cagr"] = round(
            ((income_statement_df["latest_" + column] / income_statement_df[column]) ** (1 / (n_years - 1)) - 1) * 100, 2)

    cagr_df = income_statement_df[return_columns]
    return cagr_df


def get_income_statement_ttm_rate(factors=[]) -> pd.DataFrame:
    """
    根据最新财务报表-利润表（年报、季报），获取主要指标（营业收入、利润总额、净利润、扣非净利润、每股收益）TTM增长率。

    这也是基于已公布并存库的财务报表数据处理得到的。

    ？？？todo 如果有负数，怎么处理（如果去年的基数为负，今年为正，这里返回的增长率会是负数，明显不太对）

    :param factors: 需要计算的指标，可选：[operate_income, total_profit, netprofit, deduct_parent_netprofit, basic_eps]
    :return:
    """
    query_sql = "select t.* from " \
                "(select *, rank() over(partition by code order by report_date desc) as time_rank " \
                "from `stock_income_statement` ) t where t.time_rank <= 9 order by code, report_date"
    engine = get_mysql_engine()
    income_statement_df = pd.read_sql(text(query_sql), engine.connect())
    engine.connect().close()
    if len(factors) == 0:
        columns = ["operate_income", "total_profit", "netprofit", "deduct_parent_netprofit", "basic_eps"]
    else:
        columns = factors
    return_columns = ["code", "name", "report_name"]
    for column in columns:
        return_columns.append(column + "_ttm")
        return_columns.append(column + "_ttm_rate")
        # 这一坨代码在干啥？
        # 1、先算出各指标单季度数据（除一季报外，单季度数据 = 当前财报 - 上一份财报）
        # 2、根据单季度，算出各指标最新TTM和上一年TTM数据（4个单季度数据之和）
        # 3、各指标TTM增长率 = (最新TTM / 上一年TTM) - 1
        # 不要问我这一坨魔幻的代码咋来的，不想for循环处理每只股，边查边写边调试，最后目标结果出来了，代码也变成这样了（主要得按股票分组处理）
        income_statement_df.loc[income_statement_df["report_type"] == 1, column + "_quarter"] = income_statement_df[column]
        income_statement_df.loc[income_statement_df["report_type"] != 1, column + "_quarter"] = income_statement_df[column] - income_statement_df[column].groupby(income_statement_df["code"]).shift(1)
        income_statement_df[column + "_ttm"] = income_statement_df.groupby(["code"], as_index=False)[column + "_quarter"].rolling(4).sum()[column + "_quarter"]
        income_statement_df[column + "_ttm_last"] = income_statement_df[column + "_ttm"].groupby(income_statement_df["code"]).shift(4)
        income_statement_df[column + "_ttm_rate"] = round((income_statement_df[column + "_ttm"] / income_statement_df[column + "_ttm_last"] - 1) * 100, 2)

    # 删除辅助数据行（上一份TTM数据为空）
    income_statement_df.dropna(axis=0, subset=[columns[0] + "_ttm_last"], inplace=True)
    ttm_rate_df = income_statement_df[return_columns]
    return ttm_rate_df


